package com.yfqy.app.configure;

import org.slf4j.MDC;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

@Configuration
public class ThreadPoolExecutorConfig implements DisposableBean {

    private static final int CORE_POOL_SIZE = 10;
    private static final int MAX_POOL_SIZE = 100;
    private static final long KEEP_ALIVE_TIME = 60L;
    private static final int QUEUE_CAPACITY = 1000;
    private static final String THREAD_NAME_PREFIX = "DefaultThreadPoolExecutor-";

    private ThreadPoolExecutor threadPoolExecutor;


    @Bean(name = "defaultThreadPoolExecutor")
    public ThreadPoolExecutor threadPoolExecutor() {
        this.threadPoolExecutor = new ContextAwareThreadPoolExecutor(
                CORE_POOL_SIZE,
                MAX_POOL_SIZE,
                KEEP_ALIVE_TIME,
                TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(QUEUE_CAPACITY),
                new CustomThreadFactory(THREAD_NAME_PREFIX),
                new ThreadPoolExecutor.CallerRunsPolicy()  // 添加拒绝策略:由提交任务的线程自己执行
        );
        return threadPoolExecutor;
    }

    @Override
    public void destroy() {
        shutdownThreadPoolGracefully();
    }

    private void shutdownThreadPoolGracefully() {
        if (threadPoolExecutor != null) {
            // 禁止新任务提交
            threadPoolExecutor.shutdown();
            try {
                if (!threadPoolExecutor.awaitTermination(KEEP_ALIVE_TIME, TimeUnit.SECONDS)) {
                    // 取消当前执行的任务
                    threadPoolExecutor.shutdownNow();
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                threadPoolExecutor.shutdownNow();
            }
        }
    }

    /**
     * 自定义线程工厂（命名优化）
     */
    private static class CustomThreadFactory implements ThreadFactory {
        private final String namePrefix;
        private final AtomicInteger threadCount = new AtomicInteger(1);

        CustomThreadFactory(String namePrefix) {
            this.namePrefix = namePrefix;
        }

        @Override
        public Thread newThread(Runnable r) {
            return new Thread(r, namePrefix + threadCount.getAndIncrement());
        }
    }

    private static class ContextAwareThreadPoolExecutor extends ThreadPoolExecutor {

        public ContextAwareThreadPoolExecutor(int corePoolSize, int maxPoolSize,
                                              long keepAliveTime, TimeUnit unit,
                                              BlockingQueue<Runnable> workQueue,
                                              ThreadFactory threadFactory,
                                              RejectedExecutionHandler handler) {
            super(corePoolSize, maxPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
        }

        /**
         * 接受 Runnable 类型的任务
         * 没有返回值（void 方法）
         * 无法直接获取任务执行结果
         * 任务抛出的异常会传播到未捕获异常处理器（UncaughtExceptionHandler）
         *
         * @param command the task to execute
         */
        @Override
        public void execute(Runnable command) {
            super.execute(wrap(command, MDC.getCopyOfContextMap()));
        }

        /**
         * 接受 Callable<T> 类型的任务
         * 返回 Future<T> 对象
         * 可以获取任务执行结果（通过 Future.get()）
         * 任务抛出的异常会被捕获并存储在 Future 中
         *
         * @param task the task to submit
         * @param <T>
         * @return
         */
        @Override
        public <T> Future<T> submit(Callable<T> task) {
            return super.submit(wrap(task, MDC.getCopyOfContextMap()));
        }

        /**
         * 接受 Runnable 类型的任务
         * 返回 Future<?> 对象
         * 可以检查任务是否完成（通过 Future 方法）
         * 任务正常完成时 Future.get() 返回 null
         * 任务抛出的异常会被捕获并存储在 Future 中
         *
         * @param task the task to submit
         * @return
         */
        @Override
        public Future<?> submit(Runnable task) {
            return super.submit(wrap(task, MDC.getCopyOfContextMap()));
        }

        private static Runnable wrap(Runnable runnable, Map<String, String> context) {
            return () -> {
                if (context != null && !context.isEmpty()) {
                    MDC.setContextMap(context);
                }
                try {
                    runnable.run();
                } finally {
                    MDC.clear();
                }
            };
        }

        private static <T> Callable<T> wrap(Callable<T> callable, Map<String, String> context) {
            return () -> {
                if (context != null && !context.isEmpty()) {
                    MDC.setContextMap(context);
                }
                try {
                    return callable.call();
                } finally {
                    MDC.clear();
                }
            };
        }
    }
}
